package com.eugenes.collections;

import java.util.Comparator;
import java.util.ConcurrentModificationException;
import java.util.Iterator;
import java.util.NoSuchElementException;

public class MyTreeMap<K, V> implements MyMap<K, V> {

    private SimpleEntry<K, V> root = new SimpleEntry<>(null, null, null);
    private Comparator<K> comparator = null;

    private int size;
    private int modificationCount = Integer.MIN_VALUE;

    @Override
    public void clear() {
        ++modificationCount;

        // bad boy - don't wanna help GC
        root = new SimpleEntry<>(null, null, null);
    }

    @Override
    public V put(K key, V value) {
        if (key == null) {
            throw new IllegalArgumentException("Cannot insert null key into MyTreeMap");
        }

        ++modificationCount;
        SimpleEntry<K, V> newEntry = createEntry(key, value);
        SimpleEntry<K, V> current = root;

        while (current.getKey() != null) {
            int compareResult = compare(key, current.getKey());
            if (compareResult < 0) {
                current = current.left;
            } else if (compareResult > 0) {
                current = current.right;
            } else {
                V result = current.value;
                current.value = value;
                return result;
            }
        }

        ++size;

        if (current.parent == null) {
            root = newEntry;
        } else if (current.parent.left == current) {
            setLeft(current.parent, newEntry);
        } else {
            setRight(current.parent, newEntry);
        }

        fixAdded(newEntry);
        return null;
    }

    @Override
    public V remove(Object obj) {
        SimpleEntry<K, V> entry = getEntry(obj);

        if (entry != null) {
            ++modificationCount;
            V result = entry.value;
            --size;
            removeEntry(entry);
            return result;
        }

        return null;
    }

    @Override
    public boolean containsKey(Object obj) {
        SimpleEntry<K, V> entry = getEntry(obj);
        return entry != null;
    }

    private boolean subtreeContainsValue(SimpleEntry<K, V> subtreeRoot, Object value) {
        if (subtreeRoot.getKey() == null) {
            return false;
        } else if (value == null ?
                subtreeRoot.getValue() == null : value.equals(subtreeRoot.getValue())) {
            return true;
        }

        return subtreeContainsValue(subtreeRoot.left, value)
                || subtreeContainsValue(subtreeRoot.right, value);
    }

    @Override
    public boolean containsValue(Object value) {
        return subtreeContainsValue(root, value);
    }

    @Override
    public V get(Object key) {
        SimpleEntry<K, V> entry = getEntry(key);
        return entry == null ? null : entry.getValue();
    }

    @Override
    public boolean isEmpty() {
        return 0 == size();
    }

    @Override
    public Iterator<Entry<K, V>> entryIterator() {
        return new MyTreeMapIterator();
    }

    @Override
    public int size() {
        return size;
    }

    private SimpleEntry<K, V> getEntry(Object obj) {
        SimpleEntry<K, V> current = root;

        while (current.getKey() != null) {
            int compareResult = compare(current.getKey(), obj);
            if (compareResult > 0) {
                current = current.left;
            } else if (compareResult < 0) {
                current = current.right;
            } else {
                return current;
            }
        }

        return null;
    }

    private SimpleEntry<K, V> getSmallestInSubtree(SimpleEntry<K, V> subtreeRoot) {
        if (subtreeRoot == null || subtreeRoot.getKey() == null) {
            return null;
        }
        while (subtreeRoot.getKey() != null) {
            subtreeRoot = subtreeRoot.left;
        }
        return subtreeRoot.parent;
    }

    private SimpleEntry<K, V> getLargestInSubtree(SimpleEntry<K, V> subtreeRoot) {
        if (subtreeRoot == null || subtreeRoot.getKey() == null) {
            return null;
        }
        while (subtreeRoot.getKey() != null) {
            subtreeRoot = subtreeRoot.right;
        }
        return subtreeRoot.parent;
    }

    private SimpleEntry<K, V> getSuccessor(SimpleEntry<K, V> entry) {
        if (entry.right != null && entry.right.getKey() != null) {
            return getSmallestInSubtree(entry.right);
        }

        SimpleEntry<K, V> parent = entry.parent;
        while (parent != null && parent.right == entry) {
            entry = parent;
            parent = entry.parent;
        }

        return parent;
    }

    private SimpleEntry<K, V> createEntry(K key, V value) {
        SimpleEntry<K, V> result = new SimpleEntry<>(key, value, null);
        SimpleEntry<K, V> left = new SimpleEntry<>(null, null, result);
        SimpleEntry<K, V> right = new SimpleEntry<>(null, null, result);

        result.left = left;
        result.right = right;
        left.color = SimpleEntry.Color.BLACK;
        right.color = SimpleEntry.Color.BLACK;

        return result;
    }

    private void removeEntry(SimpleEntry<K, V> entry) {
        SimpleEntry<K, V> lgstLeft = getLargestInSubtree(entry.left);

        if (lgstLeft != null) {
            entry.key = lgstLeft.key;
            entry.value = lgstLeft.value;
            entry = lgstLeft;
        }

        deleteOneChildEntry(entry);
    }

    private int compare(Object o1, Object o2) {
        return comparator == null ? ((Comparable<? super K>)o1).compareTo((K)o2) :
                comparator.compare((K)o1, (K)o2);
    }

    private static <K, V> SimpleEntry<K, V>
            getParent(SimpleEntry<K, V> entry) {
        return entry == null ? null : entry.parent;
    }

    private static <K, V> SimpleEntry<K, V>
            getGrandparent(SimpleEntry<K, V> entry) {
        return getParent(getParent(entry));
    }

    private static <K, V> SimpleEntry<K, V>
            getSibling(SimpleEntry<K, V> entry) {
        SimpleEntry<K, V> parent = getParent(entry);
        return parent == null ? null :
                parent.left == entry ? parent.right : parent.left;
    }

    private SimpleEntry<K, V> getUncle(SimpleEntry<K, V> entry) {
        SimpleEntry<K, V> parent = getParent(entry);
        SimpleEntry<K, V> grandparent = getGrandparent(entry);
        return grandparent == null ? null :
                grandparent.left == parent ? grandparent.right : grandparent.left;
    }

    private void setLeft(SimpleEntry<K, V> entry, SimpleEntry<K, V> newLeft) {
        entry.left = newLeft;
        newLeft.parent = entry;
    }

    private void setRight(SimpleEntry<K, V> entry, SimpleEntry<K, V> newRight) {
        entry.right = newRight;
        newRight.parent = entry;
    }

    private void rightRotate(SimpleEntry<K, V> a) {
        SimpleEntry<K, V> b = a.left;
        SimpleEntry<K, V> d = b.right;

        SimpleEntry<K, V> aParent = getParent(a);
        setLeft(a, d);
        setRight(b, a);
        b.parent = aParent;
        if (aParent != null) {
            if (aParent.left == a) {
                aParent.left = b;
            } else {
                aParent.right = b;
            }
        } else {
            root = b;
        }
    }

    private void leftRotate(SimpleEntry<K, V> b) {
        SimpleEntry<K, V> a = b.right;
        SimpleEntry<K, V> d = a.left;

        SimpleEntry<K, V> bParent = b.parent;
        setLeft(a, b);
        setRight(b, d);
        a.parent = bParent;
        if (bParent != null) {
            if (bParent.left == b) {
                bParent.left = a;
            } else {
                bParent.right = a;
            }
        } else {
            root = a;
        }
    }

    private void fixAdded(SimpleEntry<K, V> added) {
        if (added.parent == null) {
            added.color = SimpleEntry.Color.BLACK;
        } else {
            checkParentIsBlackAfterInsert(added);
        }
    }

    private void checkParentIsBlackAfterInsert(SimpleEntry<K, V> added) {
        SimpleEntry<K, V> parent = getParent(added);
        if (parent.color == SimpleEntry.Color.BLACK) {
            return;
        } else {
            checkUncleIsRedAfterInsert(added);
        }
    }

    private void checkUncleIsRedAfterInsert(SimpleEntry<K, V> added) {
        SimpleEntry<K, V> uncle = getUncle(added);
        if (uncle.color == SimpleEntry.Color.RED) {
            added.parent.color = SimpleEntry.Color.BLACK;
            uncle.color = SimpleEntry.Color.BLACK;
            SimpleEntry<K, V> grandparent = getGrandparent(added);
            grandparent.color = SimpleEntry.Color.RED;
            fixAdded(grandparent);
        } else {
            fixParentAndRotateAfterInsert(added);
        }
    }

    private void fixParentAndRotateAfterInsert(SimpleEntry<K, V> added) {
        SimpleEntry<K, V> grandparent = getGrandparent(added);
        if (added.parent.right == added && grandparent.left == added.parent) {
            leftRotate(added.parent);
            added = added.left;
        } else if (added.parent.left == added && grandparent.right == added.parent) {
            rightRotate(added.parent);
            added = added.right;
        }
        doRotateFixAfterInsert(added);
    }

    private void doRotateFixAfterInsert(SimpleEntry<K, V> added) {
        SimpleEntry<K, V> grandparent = getGrandparent(added);
        grandparent.color = SimpleEntry.Color.RED;
        added.parent.color = SimpleEntry.Color.BLACK;
        if (added.parent.left == added) {
            rightRotate(grandparent);
        } else {
            leftRotate(grandparent);
        }
    }

    private void replaceEntry(SimpleEntry<K, V> entry,
            SimpleEntry<K, V> replacement) {
        SimpleEntry<K, V> parent = getParent(entry);
        if (parent == null) {
            root = replacement;
        } else {
            if (parent.left == entry) {
                setLeft(parent, replacement);
            } else {
                setRight(parent, replacement);
            }
        }
        entry.key = null;
        entry.value = null;
        entry.left = entry.right = entry.parent = null;
    }

    private void deleteOneChildEntry(SimpleEntry<K, V> entry) {
        SimpleEntry<K, V> child = entry.left.key == null ?
                entry.right : entry.left;

        replaceEntry(entry, child);

        if (entry.color == SimpleEntry.Color.BLACK) {
            if (child.color == SimpleEntry.Color.RED) {
                child.color = SimpleEntry.Color.BLACK;
            } else {
                checkRootAfterDelete(child);
            }
        } else {
            child.left = child.right = null;
        }
    }

    private void checkRootAfterDelete(SimpleEntry<K, V> entry) {
        if (getParent(entry) != null) {
            doSiblingBlackAndFixAfterDelete(entry);
        }
    }

    private void doSiblingBlackAndFixAfterDelete(SimpleEntry<K, V> entry) {
        SimpleEntry<K, V> sibling = getSibling(entry);
        SimpleEntry<K, V> parent = getParent(entry);

        if (sibling.color == SimpleEntry.Color.RED) {
            parent.color = SimpleEntry.Color.RED;
            sibling.color = SimpleEntry.Color.BLACK;
            if (parent.left == entry) {
                leftRotate(parent);
            } else {
                rightRotate(parent);
            }
        }
        fixAfterDeleteSiblingIsBlack(entry);
    }

    private void fixAfterDeleteSiblingIsBlack(SimpleEntry<K, V> entry) {
        SimpleEntry<K, V> sibling = getSibling(entry);
        SimpleEntry<K, V> parent = getParent(entry);

        if (sibling.color == SimpleEntry.Color.BLACK
                && sibling.left.color == SimpleEntry.Color.BLACK
                && sibling.right.color == SimpleEntry.Color.BLACK) {
            sibling.color = SimpleEntry.Color.RED;

            if (parent.color == SimpleEntry.Color.BLACK) {
                checkRootAfterDelete(parent);
            } else {
                parent.color = SimpleEntry.Color.BLACK;
            }
        } else {
            if (sibling.color == SimpleEntry.Color.BLACK) {
                if (parent.left == entry
                        && sibling.right.color == SimpleEntry.Color.BLACK
                        && sibling.left.color == SimpleEntry.Color.RED) {
                    sibling.color = SimpleEntry.Color.RED;
                    sibling.left.color = SimpleEntry.Color.BLACK;
                    rightRotate(sibling);
                } else if (parent.right == entry
                        && sibling.left.color == SimpleEntry.Color.BLACK
                        && sibling.right.color == SimpleEntry.Color.RED) {
                    sibling.color = SimpleEntry.Color.RED;
                    sibling.right.color = SimpleEntry.Color.BLACK;
                    leftRotate(sibling);
                }
            }

            doFinalRotateAfterDelete(entry);
        }
    }

    private void doFinalRotateAfterDelete(SimpleEntry<K, V> entry) {
        SimpleEntry<K, V> sibling = getSibling(entry);
        SimpleEntry<K, V> parent = getParent(entry);

        sibling.color = parent.color;
        parent.color = SimpleEntry.Color.BLACK;

        if (parent.left == entry) {
            sibling.right.color = SimpleEntry.Color.BLACK;
            leftRotate(parent);
        } else {
            sibling.left.color = SimpleEntry.Color.BLACK;
            rightRotate(parent);
        }
    }

    private class MyTreeMapIterator implements Iterator<Entry<K, V>> {

        SimpleEntry<K, V> prevReturnedEntry;
        SimpleEntry<K, V> returnedEntry;
        int lastMod;
        boolean canRemove;

        public MyTreeMapIterator() {
            returnedEntry = null;
            lastMod = modificationCount;
            canRemove = false;
        }

        void checkForModification() {
            if (modificationCount != lastMod) {
                throw new ConcurrentModificationException("MyTreeMap changed since iterator was created.");
            }
        }

        @Override
        public boolean hasNext() {
            checkForModification();
            return returnedEntry == null ?
                    size != 0 : getSuccessor(returnedEntry) != null;
        }

        @Override
        public Entry<K, V> next() {
            checkForModification();
            if (!hasNext()) {
                throw new NoSuchElementException("No next element in this MyTreeMap instance.");
            }

            prevReturnedEntry = returnedEntry;
            returnedEntry = returnedEntry == null ?
                    getSmallestInSubtree(root) : getSuccessor(returnedEntry);
            canRemove = true;

            return returnedEntry;
        }

        @Override
        public void remove() {
            checkForModification();
            if (!canRemove) {
                throw new IllegalStateException("Next method have to be called before remove.");
            }

            canRemove = false;
            removeEntry(returnedEntry);
            returnedEntry = prevReturnedEntry;
        }

    }

    static class SimpleEntry<K, V> implements Entry<K, V> {
        static enum Color {
            BLACK, RED
        }

        K key;
        V value;

        Color color = Color.RED;
        SimpleEntry<K, V> parent;
        SimpleEntry<K, V> left = null;
        SimpleEntry<K, V> right = null;

        SimpleEntry(K key, V value, SimpleEntry<K, V> parent) {
            this.key = key;
            this.value = value;
            this.parent = parent;
        }

        @Override
        public K getKey() {
            return key;
        }

        @Override
        public V getValue() {
            return value;
        }

        @Override
        public V setValue(V value) {
            V result = this.value;
            this.value = value;
            return result;
        }
    }
}
